package grpc_client

import (
	"sync"
	"time"

	"github.com/pkg/errors"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/keepalive"
)

type ConnectionManager struct {
	sync.RWMutex
	addrs map[string]string
	conns map[string]*grpc.ClientConn // one ip, one conn
}

func NewConnectionManager() *ConnectionManager {
	return &ConnectionManager{
		addrs: make(map[string]string),
		conns: make(map[string]*grpc.ClientConn),
	}
}

func (cm *ConnectionManager) Get(appID string) (*grpc.ClientConn, error) {
	cm.RLock()
	defer cm.RUnlock()
	addr, ok := cm.addrs[appID]
	if !ok {
		return nil, errors.New("grpc app_id not exist")
	}
	conn, ok := cm.conns[addr]
	if !ok {
		return nil, errors.New("grpc conn not exist")

	}
	return conn, nil
}

func (cm *ConnectionManager) Put(appID, addr string) (*grpc.ClientConn, error) {
	cm.Lock()
	defer cm.Unlock()
	c, err := cm.newConn(addr)
	if err != nil {
		return nil, err
	}
	cm.addrs[appID] = addr
	cm.conns[addr] = c
	return c, nil
}

func (cm *ConnectionManager) Delete(appID string) {
	if addr, ok := cm.addrs[appID]; ok {
		delete(cm.conns, addr)
	}
	delete(cm.addrs, appID)
}

func (cm *ConnectionManager) newConn(addr string) (*grpc.ClientConn, error) {
	dialopts := []grpc.DialOption{
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time:                2 * time.Minute,
			Timeout:             2 * time.Minute,
			PermitWithoutStream: true,
		}),
	}

	conn, err := grpc.Dial(addr, dialopts...)
	if err != nil {
		return nil, err
	}
	return conn, nil
}
